#! /bin/bash
date >> $LOGFILE
#echo "my prolog=$0" >> $LOGFILE
#echo "my workdir=`pwd`" >> $LOGFILE

SLURM_INSTALL_PATH=/opt/gridview/slurm
# ========= job params from jobdetail==========
#job resource from job detail
#JOB_CPUIDS
DOCKER_CPUIDS=$(getFinalNodeCpus $JOBDETAIL)
DOCKER_GPUIDS=$(getFinalNodeGpus $JOBDETAIL)
DOCKER_CTRMEM=$(getFinalNodeMem $JOBDETAIL)
#echo "DOCKER_CPUIDS=$DOCKER_CPUIDS" >> $LOGFILE
echo "DOCKER_GPUIDS=$DOCKER_GPUIDS" >> $LOGFILE
#echo "DOCKER_CTRMEM=$DOCKER_CTRMEM" >> $LOGFILE

#single docker container name
DOCKER_NAME=$(getDockerName)
#echo "DOCKER_NAME=$DOCKER_NAME" >> $LOGFILE
# ===================
# private
initRoleContainerParams(){

    docker_shmsize="$DOCKER_CTRMEM"

    echo "docker_shmsize=$docker_shmsize" >> $LOGFILE
    #
    docker_framework_env=" --env JUPYTERLAB_WORKSPACES_DIR=${TASK_PATH} --cap-add=SYS_PTRACE --security-opt seccomp=unconfined "
    if [ -e ${USERHOME}/.ai_user_info/ai_proxy ];then
	 http_proxy=`cat ${USERHOME}/.ai_user_info/ai_proxy  | head -n1 | awk '{print $2}' | awk -F"'" '{print $2}'`
         docker_framework_env="$docker_framework_env --env http_proxy='$http_proxy' --env https_proxy='$http_proxy' "
    fi
    #
    docker_framework_entrypoint=" --entrypoint /bin/sh "
    #
    extra_mount_paths=`getExtraMountInfo "${extra_mount_source_path}" "${extra_mount_target_path}" "${extra_mount_target_path_permission}"`
    #
    ib_driver_mount="-v ${sharing_path}/ib_driver/:/opt/ib_driver/:ro "
    #
    motd_mount="-v ${TASK_PATH}/motd:/etc/motd"
	  #
    docker_framework_mounts="-v ${SOTHISAI_HOME}/scripts/scheduler/slurm/instance/:/opt/SothisAI/ ${extra_mount_paths} ${motd_mount} ${ib_driver_mount}"
    export TASK_TYPE=${task_type}
    if [ -z $task_command ];then
        run_command="/usr/sbin/sshd -D"
    else
        if [ ! -z $use_root ] && [ $use_root = 'true' ];then
              run_command="/usr/sbin/sshd -D &  cd ${USERHOME} ; $task_command "
        else
              run_command="/usr/sbin/sshd -D & su - ${SLURM_JOB_USER} -c 'cd ${USERHOME} ; $task_command' "
        fi
    fi

    if [ ${task_type} = "ssh" ];then
        INSTANCE_CMD=${run_command}
    elif [ ${task_type} = "jupyter" ];then
        INSTANCE_CMD=" if [ -n "${work_space_json_path}" ]; then /opt/conda/bin/jupyter lab workspaces import ${work_space_json_path}; fi ;  $run_command"
    elif [ ${task_type} = "codeserver" ];then
        prepare_codeserver_env ${USER_NAME} ${USERHOME}
        INSTANCE_CMD=${run_command}
    elif [ ${task_type} = "rstudio" ];then
        INSTANCE_CMD="chown ${SLURM_JOB_USER} /var/lib/rstudio-server/rstudio-os.sqlite* /tmp/rstudio-server/session-rpc-key /tmp/rstudio-server/secure-cookie-key ; $run_command"
    elif [ ${task_type} = "notebook_tensorboard" ] || [ ${task_type} = "instance_tensorboard" ];then
        INSTANCE_CMD="${task_command}"
    fi 
    docker_cmd_arg="cat /proc/1/environ| tr '\0' '\n' | grep -v  '^HOME=' | awk -v ex='export ' '{printf ex;print \\\$1}'> /etc/profile.d/sothisai.sh;echo \\\\\\\"export HOME=$USERHOME\\\\\\\" >> /etc/profile.d/sothisai.sh;sed -i -e 's/=/=\\\\\\\"/' -e 's/$/\\\\\\\"/' /etc/profile.d/sothisai.sh ; sed -i 's/.*UseDNS.*/UseDNS no/' /etc/ssh/sshd_config ; sed -i 's/.*env_reset.*/Defaults \!env_reset/' /etc/sudoers ; sed -i '/ai_proxy/d' /etc/bashrc /etc/bash.bashrc >/dev/null 2>&1 ; if [ -e $USERHOME/.ai_user_info/ai_proxy ];then echo 'source $USERHOME/.ai_user_info/ai_proxy' | tee -a /etc/bashrc /etc/bash.bashrc ; fi ; ${INSTANCE_CMD} "

}

# Usage: ${taskUser} ${taskImage} $dcName $dcCpuIDs ${taskMem} $dcGpuIDs $taskRole $taskx $taskNode
function startRoleContainer(){
	#echo "Arguments of startRoleContainer: $*" >> $LOGFILE
	if [ $# -lt 11 ];then
 		#echo "Arguments of startRoleContainer not valid " >> $LOGFILE
		return 1
	fi
	initRoleContainerParams
	DC_NAME=$3
	createRoleContainerParamsFile $*
	# local start
	echo "sh $SLURM_INSTALL_PATH/etc/sothisai/startRoleContainer.sh $TASK_PATH $DC_NAME $SLURM_JOB_ID" >> $LOGFILE
	sh $SLURM_INSTALL_PATH/etc/sothisai/startRoleContainer.sh $TASK_PATH $DC_NAME $SLURM_JOB_ID
}

function SingleStart()
{
    local taskUser=$SLURM_JOB_USER
    local taskImage=$IMAGENAME
    local dcName=$DOCKER_NAME
    local dcCpuIDs=$DOCKER_CPUIDS
    local taskMem=$DOCKER_CTRMEM
    local dcGpuIDs=$DOCKER_GPUIDS
    if [ "$dcGpuIDs" = "" ];then
        dcGpuIDs="-"
    fi
    local taskRole="worker"
    local taskx=0
    local taskNode=${SLURMD_NODENAME}
    dcGpuIDs_arr=($(echo ${dcGpuIDs} | sed 's/,/ /g'))
    echo " -> dcCpuIDs=$dcCpuIDs" >> $LOGFILE
    echo " -> non-distributed dcGpuIDs=$dcGpuIDs" >> $LOGFILE
    Gpu_arr=($(getCUDAGPUIDs))
    if [ $? -ne 5 ] && [ "gpu" = "${accelerator_type}" ] && [ "$dcGpuIDs" != "-" ];then
       Gpu_arr_s=${Gpu_arr[*]}
       Gpu_arr_split=($(splitCUDAGPUIDs 0 $((${#dcGpuIDs_arr[*]})) "${Gpu_arr_s}"))
       if [ "${#dcGpuIDs_arr[*]}" -eq "${#Gpu_arr_split[*]}" ] && [  "${dcGpuIDs_arr[*]}" != "${Gpu_arr_split[*]}" ];then
          dcGpuIDs=$(echo "${Gpu_arr_split[*]}" | sed 's/ /,/g')
       fi
    echo " -> new dcGpuIDs=$dcGpuIDs" >> $LOGFILE
    fi
    if [ ! -z $instance_share_dir ] && [ -f ${instance_share_dir}/shared_hosts ];then
	    instance_share_file=${instance_share_dir}/shared_hosts
    fi
    startRoleContainer ${taskUser} ${taskImage} $dcName "$dcCpuIDs" ${taskMem} "$dcGpuIDs" "$taskRole" "$taskx" "$taskNode" "$accelerator_type" "worker-$container_index" "$instance_share_file"
    if [ ! -z $instance_share_dir ] && [ -f ${instance_share_dir}/shared_hosts ];then
        dockerlistfile=`getDockerInstanceFilePath $TASK_PATH $SLURM_JOB_ID`
        container_ip_info=`cat $dockerlistfile | awk -F, '{print $13"="$1}' | awk -F= '{print $2" "$4}'`
        echo "$container_ip_info"  >> $instance_share_file
    fi


}
SingleStart
